# -*- coding: utf-8 -*-
import numpy as np
import tensorflow as tf
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt
import sys
"""
获取所有字符集
>>> import string
>>> print string.digits+string.letters
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ
"""
plt.rcParams['font.sans-serif']=['SimHei']
modelpath="./ocrdata/model"
image=Image.open("./ocrdata/inputdata.png")
text="0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
#生成训练集
def readData(image,text,showdata=False):
    interval=1#定义间隔
    length=len(text)
    width=(image.size[0]-(length+1)*interval)//length#平分长度
    # print("位图大小%s,字符集长度%s,字符平均大小(%s,%s),间隔%s" % (image.size,length,width,image.size[0],interval))
    width_with_interval=width+interval
    #分割
    boxs=[(i*(width_with_interval)+interval,interval,i*(width_with_interval)+interval+width,image.size[1]-interval) for i in range(length)]#获取分割数据
    images=np.array([np.mean(image.crop(box),-1) for box in boxs])
    #显示图像
    if (showdata):
        plt.figure()
        for i in range(length):
            ax=plt.subplot(1,length,i+1)
            ax.axis('off')
            ax.set(title=text[i])
            plt.imshow(images[i])
        plt.show()
    return images,np.array([i for i in range(length)])
#显示训练数据
def plot_history(history):
    hist = pd.DataFrame(history.history)
    hist['epoch'] = history.epoch
    plt.figure()
    plt.xlabel('Epoch[纪元]')
    plt.ylabel('accuracy [精确率]')
    plt.plot(hist['epoch'], hist['accuracy'],
             label='Train Error')
    plt.show()
#训练
def tran(images,texts,savemodel=False,showcurve=False):
    model = tf.keras.models.Sequential([
      tf.keras.layers.Flatten(input_shape=(images.shape[1], images.shape[2])), #展平层
    #   tf.keras.layers.Dense(128, activation='relu'), #连接层
      tf.keras.layers.Dense(len(texts), activation='softmax') #归并层
    ])

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # 显示模型的结构
    model.summary()
    
    history=model.fit(images, texts, epochs=120)#120代左右时能100%识别
    if showcurve:#如果显示曲线
        plot_history(history)
    model.evaluate(images,  texts, verbose=2)
    if (savemodel):
        model.save(modelpath)
    return model

#预测矩阵
def predict(image,model=None):
    #热加载模型
    if (model == None):
        global tfmodel
        if 'tfmodel' in  globals().keys():
            model=tfmodel
        else:
            tfmodel=model=tf.keras.models.load_model(modelpath)
    img = (np.expand_dims(image,0))
    # img = np.where(img[...,:] < 90, 0, 255)
    # from app import debugshow
    # debugshow()
    # print(img)
    

    index=np.argmax(model.predict(img)[0])
    return text[index]

#预测图片
def predictimg(image,model=None):
    pass
    index=1
    return text[index]
#预测灵活处理
def predictflex(image,model=None):
    print(image.size)
    # image=np.array(image)[0:12,0:7]
    # image=np.array(image)[0:12,0:7]
    image=tf.image.resize(image, [12, 7])
    image=np.mean(image,-1)
    # print(image)
    # from app import debugshow
    # debugshow(image)
    return predict(image,model)

if __name__ == '__main__':

    if len(sys.argv)>0:
        arg=sys.argv[1]
        if arg=='read':
            # 获取数据
            readData(image,text,True)
        elif arg=='tran':
            # 训练得到模型
            model=tran(*readData(image,text),savemodel=True,showcurve=True)
        elif arg=='predict':
            # 预测
            i=0;#第0个字符
            if (sys.argv[2]):
                i=text.find(sys.argv[2])
            img=readData(image,text)[0][i]
            # from app import debugshow
            # debugshow(img)
            # print(img)
            print("字符%s预测结果为%s" % (text[i],predict(img)))